"""
File   : json_schema.py
Desc   : langchain 结构化输出: JSON Schema
         注意: langchain-core>=0.2.26
Date   : 2024/12/03
Author : Tianyu Chen
https://python.langchain.com/docs/how_to/structured_output/#pydantic-class
"""

from langchain_openai import ChatOpenAI

# Model
llm = ChatOpenAI(model="gpt-4o-mini")

# JSON Schema
json_schema = {
    "title": "joke",
    "description": "Joke to tell user.",
    "type": "object",
    "properties": {
        "setup": {
            "type": "string",
            "description": "The setup of the joke",
        },
        "punchline": {
            "type": "string",
            "description": "The punchline to the joke",
        },
        "rating": {
            "type": "integer",
            "description": "How funny the joke is, from 1 to 10",
            "default": None,
        },
    },
    "required": ["setup", "punchline", "rating"],
}

# Runnable
structured_llm = llm.with_structured_output(json_schema)

# invoke
a_joke = structured_llm.invoke("Tell me a joke about cats")
print(type(a_joke))  # <class 'dict'>
print(a_joke)  # {'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the mouse!', 'rating': 7}
print("setup:", a_joke["setup"])  # setup: Why was the cat sitting on the computer?
print("punchline:", a_joke["punchline"])  # punchline: Because it wanted to keep an eye on the mouse!
print("rating:", a_joke["rating"])  # rating: 7
